"""
Database Manager for ArXiv Scraper Service

PostgreSQL database integration with connection pooling,
paper storage, deduplication, and efficient batch operations.
"""

import logging
import hashlib
from datetime import datetime, timezone
from typing import List, Dict, Any, Tuple, Optional
import json

import psycopg2
from psycopg2 import pool, sql, extras
from psycopg2.extras import RealDictCursor, execute_batch

from utils.exceptions import DatabaseError


class DatabaseManager:
    """
    PostgreSQL database manager for paper storage and retrieval.
    
    Features:
    - Connection pooling for high performance
    - Batch operations for efficiency
    - Smart deduplication
    - Paper metadata normalization
    - Full-text search preparation
    """
    
    def __init__(self, connection_string: str, pool_size: int = 5,
                 max_overflow: int = 10):
        """Initialize database manager.
        
        Args:
            connection_string: PostgreSQL connection string
            pool_size: Initial connection pool size
            max_overflow: Maximum pool overflow
        """
        self.connection_string = connection_string
        self.pool_size = pool_size
        self.max_overflow = max_overflow
        
        self.logger = logging.getLogger(__name__)
        
        # Initialize connection pool
        self._initialize_connection_pool()
        
        # Initialize database schema
        self._initialize_schema()
        
        self.logger.info("Database manager initialized")
    
    def _initialize_connection_pool(self):
        """Initialize PostgreSQL connection pool."""
        try:
            self.connection_pool = pool.ThreadedConnectionPool(
                minconn=1,
                maxconn=self.pool_size + self.max_overflow,
                dsn=self.connection_string
            )
            
            self.logger.info(f"Connection pool initialized: "
                           f"min=1, max={self.pool_size + self.max_overflow}")
            
        except Exception as e:
            raise DatabaseError(f"Failed to initialize connection pool: {e}")
    
    def _get_connection(self):
        """Get connection from pool."""
        try:
            return self.connection_pool.getconn()
        except Exception as e:
            raise DatabaseError(f"Failed to get database connection: {e}")
    
    def _put_connection(self, conn):
        """Return connection to pool."""
        try:
            self.connection_pool.putconn(conn)
        except Exception as e:
            self.logger.error(f"Error returning connection to pool: {e}")
    
    def _initialize_schema(self):
        """Initialize database schema if it doesn't exist."""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                # Create papers table
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS papers (
                        id BIGSERIAL PRIMARY KEY,
                        arxiv_id VARCHAR(50) UNIQUE NOT NULL,
                        oai_identifier VARCHAR(255) UNIQUE NOT NULL,
                        title TEXT NOT NULL,
                        authors TEXT[] NOT NULL,
                        abstract TEXT,
                        subject_categories TEXT[],
                        primary_subject VARCHAR(20),
                        doi VARCHAR(100),
                        journal_ref TEXT,
                        comments TEXT,
                        date_created TIMESTAMP WITH TIME ZONE NOT NULL,
                        date_updated TIMESTAMP WITH TIME ZONE,
                        date_submitted TIMESTAMP WITH TIME ZONE NOT NULL,
                        version INTEGER DEFAULT 1,
                        pdf_url VARCHAR(500),
                        source_url VARCHAR(500),
                        metadata JSONB,
                        content_hash VARCHAR(64) NOT NULL,
                        processing_status VARCHAR(20) DEFAULT 'pending',
                        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                    );
                """)
                
                # Create indexes for performance
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_arxiv_id ON papers(arxiv_id);
                """)
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_date_submitted ON papers(date_submitted);
                """)
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_primary_subject ON papers(primary_subject);
                """)
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_content_hash ON papers(content_hash);
                """)
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_processing_status ON papers(processing_status);
                """)
                
                # Full-text search index on title and abstract
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_papers_fulltext 
                    ON papers USING gin(to_tsvector('english', title || ' ' || COALESCE(abstract, '')));
                """)
                
                # Create scraper_state table for persistence
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS scraper_state (
                        key VARCHAR(100) PRIMARY KEY,
                        value JSONB NOT NULL,
                        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                    );
                """)
                
                # Create scraper_metrics table
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS scraper_metrics (
                        id BIGSERIAL PRIMARY KEY,
                        metric_name VARCHAR(100) NOT NULL,
                        metric_value NUMERIC NOT NULL,
                        metric_metadata JSONB,
                        timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                    );
                """)
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_scraper_metrics_name_timestamp 
                    ON scraper_metrics(metric_name, timestamp DESC);
                """)
                
                conn.commit()
                self.logger.info("Database schema initialized successfully")
                
        except Exception as e:
            conn.rollback()
            raise DatabaseError(f"Failed to initialize database schema: {e}")
        finally:
            self._put_connection(conn)
    
    def store_papers_batch(self, papers: List[Dict[str, Any]]) -> Tuple[int, int]:
        """
        Store batch of papers with deduplication.
        
        Args:
            papers: List of paper dictionaries
            
        Returns:
            Tuple of (papers_stored, papers_skipped)
        """
        if not papers:
            return 0, 0
        
        conn = self._get_connection()
        papers_stored = 0
        papers_skipped = 0
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                # Prepare batch data
                batch_data = []
                content_hashes = set()
                
                for paper in papers:
                    # Generate content hash for deduplication
                    content_hash = self._generate_content_hash(paper)
                    
                    # Skip if duplicate within batch
                    if content_hash in content_hashes:
                        papers_skipped += 1
                        continue
                    
                    content_hashes.add(content_hash)
                    
                    batch_data.append(self._prepare_paper_data(paper, content_hash))
                
                if not batch_data:
                    return 0, len(papers)
                
                # Check for existing papers by content hash
                existing_hashes = self._get_existing_content_hashes(
                    cursor, list(content_hashes)
                )
                
                # Filter out existing papers
                new_papers = []
                for paper_data in batch_data:
                    if paper_data['content_hash'] not in existing_hashes:
                        new_papers.append(paper_data)
                    else:
                        papers_skipped += 1
                
                # Insert new papers
                if new_papers:
                    insert_query = """
                        INSERT INTO papers (
                            arxiv_id, oai_identifier, title, authors, abstract,
                            subject_categories, primary_subject, doi, journal_ref,
                            comments, date_created, date_updated, date_submitted,
                            version, pdf_url, source_url, metadata, content_hash,
                            processing_status
                        ) VALUES (
                            %(arxiv_id)s, %(oai_identifier)s, %(title)s, %(authors)s, %(abstract)s,
                            %(subject_categories)s, %(primary_subject)s, %(doi)s, %(journal_ref)s,
                            %(comments)s, %(date_created)s, %(date_updated)s, %(date_submitted)s,
                            %(version)s, %(pdf_url)s, %(source_url)s, %(metadata)s, %(content_hash)s,
                            %(processing_status)s
                        )
                        ON CONFLICT (arxiv_id) DO UPDATE SET
                            title = EXCLUDED.title,
                            authors = EXCLUDED.authors,
                            abstract = EXCLUDED.abstract,
                            date_updated = EXCLUDED.date_updated,
                            version = EXCLUDED.version,
                            metadata = EXCLUDED.metadata,
                            content_hash = EXCLUDED.content_hash,
                            updated_at = NOW()
                        WHERE papers.content_hash != EXCLUDED.content_hash;
                    """
                    
                    execute_batch(cursor, insert_query, new_papers, page_size=100)
                    papers_stored = cursor.rowcount
                
                conn.commit()
                
                self.logger.info(f"Batch processed: {papers_stored} stored, "
                               f"{papers_skipped} skipped (duplicates)")
                
                return papers_stored, papers_skipped
                
        except Exception as e:
            conn.rollback()
            self.logger.error(f"Error storing paper batch: {e}")
            raise DatabaseError(f"Failed to store paper batch: {e}")
        finally:
            self._put_connection(conn)
    
    def _prepare_paper_data(self, paper: Dict[str, Any], 
                          content_hash: str) -> Dict[str, Any]:
        """Prepare paper data for database insertion."""
        return {
            'arxiv_id': paper.get('arxiv_id', ''),
            'oai_identifier': paper.get('oai_identifier', ''),
            'title': paper.get('title', ''),
            'authors': paper.get('authors', []),
            'abstract': paper.get('abstract', ''),
            'subject_categories': paper.get('subject_categories', []),
            'primary_subject': paper.get('primary_subject', ''),
            'doi': paper.get('doi'),
            'journal_ref': paper.get('journal_ref'),
            'comments': paper.get('comments'),
            'date_created': paper.get('date_created'),
            'date_updated': paper.get('date_updated'),
            'date_submitted': paper.get('date_submitted'),
            'version': paper.get('version', 1),
            'pdf_url': paper.get('pdf_url'),
            'source_url': paper.get('source_url'),
            'metadata': json.dumps(paper.get('metadata', {})),
            'content_hash': content_hash,
            'processing_status': 'pending'
        }
    
    def _generate_content_hash(self, paper: Dict[str, Any]) -> str:
        """Generate content hash for deduplication."""
        # Use title, authors, and abstract for content hash
        content_parts = [
            paper.get('title', ''),
            str(paper.get('authors', [])),
            paper.get('abstract', '')[:1000]  # First 1000 chars of abstract
        ]
        
        content_string = '|'.join(content_parts).lower()
        return hashlib.sha256(content_string.encode('utf-8')).hexdigest()
    
    def _get_existing_content_hashes(self, cursor, hashes: List[str]) -> set:
        """Get existing content hashes from database."""
        if not hashes:
            return set()
        
        cursor.execute(
            "SELECT content_hash FROM papers WHERE content_hash = ANY(%s)",
            (hashes,)
        )
        
        return {row['content_hash'] for row in cursor.fetchall()}
    
    def get_paper_count(self) -> int:
        """Get total number of papers in database."""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute("SELECT COUNT(*) FROM papers")
                return cursor.fetchone()[0]
        except Exception as e:
            raise DatabaseError(f"Failed to get paper count: {e}")
        finally:
            self._put_connection(conn)
    
    def get_latest_papers(self, limit: int = 10) -> List[Dict[str, Any]]:
        """Get latest papers by submission date."""
        conn = self._get_connection()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("""
                    SELECT arxiv_id, title, authors, date_submitted, primary_subject
                    FROM papers 
                    ORDER BY date_submitted DESC 
                    LIMIT %s
                """, (limit,))
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            raise DatabaseError(f"Failed to get latest papers: {e}")
        finally:
            self._put_connection(conn)
    
    def get_papers_by_subject(self, subject: str, 
                            limit: int = 100) -> List[Dict[str, Any]]:
        """Get papers by subject category."""
        conn = self._get_connection()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("""
                    SELECT arxiv_id, title, authors, abstract, date_submitted
                    FROM papers 
                    WHERE primary_subject = %s OR %s = ANY(subject_categories)
                    ORDER BY date_submitted DESC 
                    LIMIT %s
                """, (subject, subject, limit))
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            raise DatabaseError(f"Failed to get papers by subject: {e}")
        finally:
            self._put_connection(conn)
    
    def search_papers(self, query: str, limit: int = 50) -> List[Dict[str, Any]]:
        """Full-text search in papers."""
        conn = self._get_connection()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("""
                    SELECT arxiv_id, title, authors, abstract, date_submitted,
                           ts_rank(to_tsvector('english', title || ' ' || COALESCE(abstract, '')), 
                                   plainto_tsquery('english', %s)) as rank
                    FROM papers 
                    WHERE to_tsvector('english', title || ' ' || COALESCE(abstract, '')) 
                          @@ plainto_tsquery('english', %s)
                    ORDER BY rank DESC, date_submitted DESC
                    LIMIT %s
                """, (query, query, limit))
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            raise DatabaseError(f"Failed to search papers: {e}")
        finally:
            self._put_connection(conn)
    
    def get_database_stats(self) -> Dict[str, Any]:
        """Get database statistics."""
        conn = self._get_connection()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                stats = {}
                
                # Total papers
                cursor.execute("SELECT COUNT(*) as total_papers FROM papers")
                stats['total_papers'] = cursor.fetchone()['total_papers']
                
                # Papers by processing status
                cursor.execute("""
                    SELECT processing_status, COUNT(*) as count 
                    FROM papers 
                    GROUP BY processing_status
                """)
                stats['by_status'] = {row['processing_status']: row['count'] 
                                    for row in cursor.fetchall()}
                
                # Papers by subject (top 10)
                cursor.execute("""
                    SELECT primary_subject, COUNT(*) as count 
                    FROM papers 
                    WHERE primary_subject IS NOT NULL
                    GROUP BY primary_subject 
                    ORDER BY count DESC 
                    LIMIT 10
                """)
                stats['top_subjects'] = [dict(row) for row in cursor.fetchall()]
                
                # Recent papers (last 30 days)
                cursor.execute("""
                    SELECT COUNT(*) as recent_papers 
                    FROM papers 
                    WHERE date_submitted >= NOW() - INTERVAL '30 days'
                """)
                stats['recent_papers'] = cursor.fetchone()['recent_papers']
                
                return stats
                
        except Exception as e:
            raise DatabaseError(f"Failed to get database stats: {e}")
        finally:
            self._put_connection(conn)
    
    def store_state(self, key: str, value: Any):
        """Store scraper state."""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute("""
                    INSERT INTO scraper_state (key, value, updated_at) 
                    VALUES (%s, %s, NOW())
                    ON CONFLICT (key) DO UPDATE SET 
                        value = EXCLUDED.value, 
                        updated_at = NOW()
                """, (key, json.dumps(value)))
                conn.commit()
        except Exception as e:
            raise DatabaseError(f"Failed to store state: {e}")
        finally:
            self._put_connection(conn)
    
    def get_state(self, key: str) -> Any:
        """Get scraper state."""
        conn = self._get_connection()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("SELECT value FROM scraper_state WHERE key = %s", (key,))
                result = cursor.fetchone()
                return json.loads(result['value']) if result else None
        except Exception as e:
            raise DatabaseError(f"Failed to get state: {e}")
        finally:
            self._put_connection(conn)
    
    def health_check(self) -> bool:
        """Perform database health check."""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("SELECT 1")
                cursor.fetchone()
            self._put_connection(conn)
            return True
        except Exception as e:
            self.logger.error(f"Database health check failed: {e}")
            return False
    
    def is_connected(self) -> bool:
        """Check if database connection pool is available."""
        return self.connection_pool is not None
    
    def close(self):
        """Close all database connections."""
        if hasattr(self, 'connection_pool') and self.connection_pool:
            self.connection_pool.closeall()
            self.logger.info("Database connections closed")